package cn.tpshion.vm.handler;

import cn.hutool.core.util.IdUtil;
import cn.tpshion.vm.base.config.NodeConfig;
import cn.tpshion.vm.common.constant.MsgType;
import cn.tpshion.vm.common.exception.BizException;
import cn.tpshion.vm.domain.support.MsgData;
import cn.tpshion.vm.domain.support.NodeContext;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.MediaType;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.File;
import java.util.Arrays;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class SSEHandle {

    private static final Map<String, Map<String, SseEmitter>> taskSSEMap = new ConcurrentHashMap<>();

    @Resource
    private ThreadPoolTaskExecutor msgTaskExecutor;

    @Resource
    private NodeConfig nodeConfig;

    public SseEmitter connect(String taskId, String id) {
        if(taskSSEMap.containsKey(taskId)){
            Map<String, SseEmitter> map = taskSSEMap.get(taskId);
            if(map.containsKey(id)){
                return map.get(id);
            }else{
                return createSSE(map);
            }
        }else{
            Map<String, SseEmitter> map = new ConcurrentHashMap<>();
            taskSSEMap.put(taskId, map);
            return createSSE(map);
        }
    }

    private SseEmitter createSSE(Map<String, SseEmitter> map){
        String id = IdUtil.simpleUUID();
        try {
            // 设置超时时间，0表示不过期。默认30秒
            SseEmitter sseEmitter = new SseEmitter(0L);
            MsgData msgData = new MsgData(MsgType.SUCCESS, nodeConfig.getLogo());
            sseEmitter.send(SseEmitter.event().id(id).data(msgData, MediaType.APPLICATION_JSON));
            // 注册回调
            sseEmitter.onCompletion(() -> {
                log.info("connect Completion id:{}", id);
            });
            sseEmitter.onError(throwable -> {
                log.error("connect error:{}", throwable.getMessage(), throwable);
                map.remove(id);
                throw new BizException("连接失败");
            });
            sseEmitter.onTimeout(() -> {
                log.error("connect timeout");
                map.remove(id);
                throw new BizException("连接超时");
            });
            map.put(id, sseEmitter);
            log.info("创建sse连接完成，当前id：{}", id);
            return sseEmitter;
        } catch (Exception e) {
            log.info("创建sse连接异常，当前id：{}", id);
            throw new BizException("创建sse连接异常");
        }
    }

    public void sendMessage(NodeContext nodeContext, MsgData msg){
        if(Objects.isNull(nodeContext) || Objects.isNull(msg)){
            return;
        }

        String taskId = nodeContext.getTaskId();
        Long batchId = nodeContext.getBatchId();

        if(!taskSSEMap.containsKey(taskId)){
            return;
        }

        Map<String, SseEmitter> map = taskSSEMap.get(taskId);
        map.forEach((id, sseEmitter) -> {
            try {
                sseEmitter.send(SseEmitter.event().id(id).data(msg, MediaType.APPLICATION_JSON));
            }catch (Exception e) {
                log.info("taskId:{} id:{} 推送异常:{}", taskId, id, e.getMessage());
                map.remove(id);
                sseEmitter.complete();
            }
        });

        if(!MsgType.FLOW.equals(msg.getType()) && !MsgType.NODE.getType().equals(msg.getType())){
            msgTaskExecutor.execute(() -> {
                if(StringUtils.isBlank(msg.getMsg())){
                    return;
                }
                try {
                    File file = new File(nodeConfig.getWorkspace() + File.separator + taskId + "/log");
                    if(!file.exists() || !file.isDirectory()){
                        file.mkdirs();
                    }
                    File logFile = new File(file.getPath() + File.separator + batchId + ".log");
                    FileUtils.writeLines(logFile, Arrays.asList(msg.getTime() + " " + msg.getMsg()), true);
                }catch (Exception e){
                    log.error("log write error:{}", e.getMessage());
                }
            });
        }
    }

    public void closeConnect(String taskId, String id){
        if(!taskSSEMap.containsKey(taskId)){
            return;
        }
        Map<String, SseEmitter> map = taskSSEMap.get(taskId);
        if(!map.containsKey(id)){
            return;
        }

        SseEmitter sseEmitter = map.remove(id);
        if(Objects.nonNull(sseEmitter)){
            sseEmitter.complete();
        }
    }
}
